import torch
import torch.nn as nn

from .base import BaseHead
from ..registry import HEADS
from torchmetrics.functional import accuracy


@HEADS.register()
class STGCNHead(BaseHead):
    """
    Head for ST-GCN model.
    Args:
        in_channels: int, input feature channels. Default: 256.
        num_classes: int, number classes. Default: 10.
    """

    def __init__(self, in_channels=256, num_classes=10, loss_cfg=dict(name='CrossEntropyLoss'), **kwargs):
        super().__init__(num_classes, in_channels, loss_cfg, **kwargs)
        self.fcn = nn.Conv2d(in_channels=in_channels,
                             out_channels=num_classes,
                             kernel_size=1)

    def forward(self, x):
        """Define how the head is going to run.
        """
        x = self.fcn(x)
        x = torch.reshape(x, (x.shape[0], -1))  # N,C,1,1 --> N,C

        return x

